# STAT 415/615 Regression (M. Baron)

# Python Lab 6. Simultaneous Estimation

# (Sec. 4.1 and 30.2 [in Applied Linear Statistical Models book])

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import statsmodels.api as sm
from scipy import stats

# Load the mtcars data

mtcars = sm.datasets.get_rdataset("mtcars").data

# Standard regression

X = sm.add_constant(mtcars["wt"])
y = mtcars["mpg"]
reg = sm.OLS(y, X).fit()

# Save coefficients

b = reg.params

print(b)

const    37.285126
wt       -5.344472
dtype: float64

# 97.5% confidence intervals (with Bonferroni adjustment)

CI = reg.conf_int(alpha=0.025)

print(CI)

               0          1
const  32.854747  41.715505
wt     -6.663705  -4.025238

# 95% confidence ellipse for β = (β0, β1)

# Calculate the covariance matrix of the estimated coefficients

cov = reg.cov_params()

# Critical value for a 95% confidence ellipse

chi2_critical = stats.chi2.ppf(0.95, 2)

# Create points for the confidence ellipse

theta = np.linspace(0, 2 * np.pi, 200)
circle = np.array([np.cos(theta), np.sin(theta)])

# Transform the circle using the covariance matrix

eigenvalues, eigenvectors = np.linalg.eigh(cov)
ellipse_points = (
eigenvectors
@ np.diag(np.sqrt(eigenvalues * chi2_critical))
@ circle
)

# Center the ellipse at the estimated coefficients

ellipse_points[0, :] += b.iloc[0]
ellipse_points[1, :] += b.iloc[1]

# Bonferroni confidence intervals

b0_lower = CI.iloc[0, 0]
b0_upper = CI.iloc[0, 1]

b1_lower = CI.iloc[1, 0]
b1_upper = CI.iloc[1, 1]

# Plot the confidence ellipse

plt.figure(figsize=(8, 6))

plt.plot(ellipse_points[0, :], ellipse_points[1, :], linewidth=3, label="95% confidence ellipse")

# Point estimator of vector β = (β0, β1)

plt.plot(b.iloc[0], b.iloc[1], "o")

# Draw the Bonferroni confidence rectangle

plt.plot([b0_lower, b0_upper], [b1_lower, b1_lower], linewidth=2)
plt.plot([b0_lower, b0_upper], [b1_upper, b1_upper], linewidth=2)
plt.plot([b0_lower, b0_lower], [b1_lower, b1_upper], linewidth=2)
plt.plot([b0_upper, b0_upper], [b1_lower, b1_upper], linewidth=2)

plt.xlabel(r"$\beta_0$")
plt.ylabel(r"$\beta_1$")
plt.title("Confidence ellipse and Bonferroni confidence intervals")
plt.show()